AllGather

将各输入分片的数据进行 全量收集（All-Gather）， 并在输出中按指定维度进行拼接复制。 该算子通常用于并行/多核场景中，将每个 rank 的本地数据 收集成一个完整的输出张量。

算子行为可概括为两步： 1. 将 input 中每个 input_rank 对应的数据块拷贝到 output 的首段连续区域； 2. 将首段完整数据复制到其余 output_rank - 1 个位置，形成完整聚合结果。

输入：

  • input - 输入数据地址，按 [input_rank, data_size] 形式连续存储。

  • input_rank - 输入 rank 数量（通常对应参与 AllGather 的并行单元数）。

  • output_rank - 输出 rank 数量（通常等于或大于 input_rank）。

  • data_size - 单个 rank 对应的数据元素个数。

  • core_mask - 核掩码（仅适用于共享存储版本）。

输出：

  • output - AllGather 结果输出地址， 数据按 [output_rank, input_rank, data_size] 的逻辑顺序存储。

支持平台：

FT78NE MT7004

备注

  • FT78NE 支持 fp32、fp64、int8、int16、int32、cplx64、cplx128 类型

  • MT7004 支持 fp16、fp32、int16、int32、cplx64 类型

  • 当前实现基于 DMA进行数据搬运

  • 输入与输出地址需保证足够的连续空间以容纳聚合结果

共享存储版本:

void fp_allgather_s(float *input, float *output, int input_rank, int output_rank, int data_size, int core_mask)
void dp_allgather_s(double *input, double *output, int input_rank, int output_rank, int data_size, int core_mask)
void i8_allgather_s(int8_t *input, int8_t *output, int input_rank, int output_rank, int data_size, int core_mask)
void i16_allgather_s(int16_t *input, int16_t *output, int input_rank, int output_rank, int data_size, int core_mask)
void i32_allgather_s(int32_t *input, int32_t *output, int input_rank, int output_rank, int data_size, int core_mask)
void c64_allgather_s(cplx64 *input, cplx64 *output, int input_rank, int output_rank, int data_size, int core_mask)
void c128_allgather_s(cplx128 *input, cplx128 *output, int input_rank, int output_rank, int data_size, int core_mask)

C调用示例：

 1// FT78NE 示例
 2#include <stdio.h>
 3#include <allgather.h>
 4
 5int main(int argc, char* argv[]) {
 6    float *input = (float *)0xA0000000;
 7    float *output = (float *)0xC0000000;
 8
 9    int input_rank = 4;
10    int output_rank = 4;
11    int data_size = 256;
12    int core_mask = 0xff;
13
14    fp_allgather_s(input, output,
15                   input_rank, output_rank,
16                   data_size, core_mask);
17    return 0;
18}

私有存储版本:

void fp_allgather_p(float *input, float *output, int input_rank, int output_rank, int data_size)
void dp_allgather_p(double *input, double *output, int input_rank, int output_rank, int data_size)
void i8_allgather_p(int8_t *input, int8_t *output, int input_rank, int output_rank, int data_size)
void i16_allgather_p(int16_t *input, int16_t *output, int input_rank, int output_rank, int data_size)
void i32_allgather_p(int32_t *input, int32_t *output, int input_rank, int output_rank, int data_size)
void c64_allgather_p(cplx64 *input, cplx64 *output, int input_rank, int output_rank, int data_size)
void c128_allgather_p(cplx128 *input, cplx128 *output, int input_rank, int output_rank, int data_size)

C调用示例：

 1// FT78NE 示例
 2#include <stdio.h>
 3#include <allgather.h>
 4
 5int main(int argc, char* argv[]) {
 6    float *input = (float *)0x10810000;   // L2 空间
 7    float *output = (float *)0x10820000;
 8
 9    int input_rank = 4;
10    int output_rank = 4;
11    int data_size = 256;
12
13    fp_allgather_p(input, output,
14                   input_rank, output_rank,
15                   data_size);
16    return 0;
17}